'''
Created on Jun 5, 2012

@author: nzhao
'''
import numpy as np
import matplotlib.pyplot as plt

def gso(base):
    #print base.shape
    new_base=np.zeros_like((base), dtype=float)
    new_base[0]=base[0]
    for i in range(1, len(base)):
        new_base[i]=base[i]
        for k in range(i):
            new_base[i]=new_base[i]-proj(new_base[k], base[i])
        new_base[i]=new_base[i]/np.sqrt(np.dot(new_base[i], new_base[i]))
    new_base[0]=new_base[0]/np.sqrt(np.dot(new_base[0], new_base[0]))
    
    return new_base

def proj(u, v):
    pro=np.dot(v, u)/np.dot(u, u)*u
    return pro

class super_block():
    def __init__(self, chainLen=8, truncate=8):
        self.chain_length = chainLen
        self.site_num = chainLen
        self.m = truncate
        self.zeroMat= np.zeros((self.m, self.m))        
        
        self.hamiltonian = 2.0 * np.identity(self.chain_length) \
                          -1.0 * np.diag(np.ones(self.chain_length-1),  1) \
                          -1.0 * np.diag(np.ones(self.chain_length-1), -1)

        self.tunelling = np.zeros( (self.chain_length, self.chain_length) )
        self.tunelling[-1, 0] = -1.0
       
        u, v = self.diagonalize(self.hamiltonian)
        self.RG_basis = v[:, :self.m]

        self.RG_hamiltonian = np.dot( self.RG_basis.conj().T, np.dot(self.hamiltonian, self.RG_basis) )
        self.RG_tunelling   = np.dot( self.RG_basis.conj().T, np.dot(self.tunelling,   self.RG_basis) )

        self.energy   = u[:self.m]                
        self.wavefunction = self.RG_basis.T

    def diagonalize(self, mat):
        [u,v] = np.linalg.eigh(mat)
        idx = u.argsort()
        return [ u[idx], v[:, idx] ]

    def super_block_hamiltonian(self):
        self.sb_hamiltonian = np.kron(np.eye(self.p), self.RG_hamiltonian) \
            +np.kron(np.eye(self.p, self.p, 1)+np.eye(self.p, self.p, 1-self.p), self.RG_tunelling)\
            +np.kron(np.eye(self.p, self.p, -1)+np.eye(self.p, self.p, self.p-1), self.RG_tunelling.conj().T)
 

    def block_renormalize(self, sb_size=10):
        self.site_num = self.site_num * 2
        self.p = sb_size
        
        self.hamiltonian = np.r_[ 
                np.c_[self.RG_hamiltonian,        self.RG_tunelling],
                np.c_[self.RG_tunelling.conj().T, self.RG_hamiltonian] ]
        
        self.tunelling =   np.r_[ 
                np.c_[self.zeroMat,      self.zeroMat],
                np.c_[self.RG_tunelling, self.zeroMat] ]
        
        self.super_block_hamiltonian()
        u, v = self.diagonalize(self.sb_hamiltonian)
        self.RG_basis = gso(v[:2*self.m, :self.m].T).T
        
        self.RG_hamiltonian = np.dot(self.RG_basis.conj().T, np.dot(self.hamiltonian, self.RG_basis) )
        self.RG_tunelling   = np.dot(self.RG_basis.conj().T, np.dot(self.tunelling,   self.RG_basis) )
        
        u1, v1 = self.diagonalize(self.RG_hamiltonian)
        self.energy = u1[:self.m]
        self.wavefunction = np.dot(self.wavefunction.T, v1).T

    
b=super_block()
for i in range(8):
    b.block_renormalize()
    print i, b.energy
